package org.wsocket.protocol;

import java.io.IOException;
import java.io.InputStream;
import java.io.PrintStream;
import java.io.UnsupportedEncodingException;
import java.math.BigInteger;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import java.nio.charset.Charset;
import java.security.MessageDigest;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.log4j.Logger;
import org.wsocket.WSocket;
import org.wsocket.connector.Connector;

/**
 * This class used to process client socket request/response
 * 
 * @author  JoJo, Alvin
 * 
 */
public class SocketHandler {
    public static String PATHARG_SEPARATOR = ";";
    public static String ARGARG_SEPARATOR = ",";
    public static String KEYVAL_SEPARATOR = "=";
    public static String KEY_PARAM = "key";
    public static int BUFFER_SIZE = 1024 * 1024;
    public static String CHARSET = "UTF-8";

    private static Logger log=Logger.getLogger(SocketHandler.class);
    /**
     * read request from socket client
     * 
     */
    public static String processRequest(SelectionKey sk,
    		ConcurrentHashMap<String, Connector> connectors) {
        SocketChannel socketChannel = (SocketChannel) sk.channel();
        if (!socketChannel.isOpen()) {
            return null;
        }
        ByteBuffer buff = ByteBuffer.allocate(BUFFER_SIZE);
        String clientMessage = null;
        Connector t = (Connector) sk.attachment();

        try {
            while (true) {
                buff.clear();
                int r = socketChannel.read(buff);
                if (r <= 0) {
                    break;
                } else {
                    sk.interestOps(SelectionKey.OP_READ);
                }
                buff.flip();
                if (clientMessage == null) {
                    clientMessage = "" + Charset.forName(CHARSET).decode(buff);
                } else {
                    clientMessage += Charset.forName(CHARSET).decode(buff);
                }

            }
            if (null == clientMessage) {
                throw new IOException("log out.");
            }
            if ("".equals(clientMessage)) {
                return null;
            }

            // cut off potential starting 0x00 and 0xff characters
            byte[] ba = clientMessage.getBytes(CHARSET);
            int i = 0;
            while (i < ba.length && ba[i] != 0) {
                i++;
            }
            if (i < (ba.length - 3)) {
                i++;
                clientMessage = new String(ba, i, ba.length - 3 - i, CHARSET);
            } else {
                return null;
            }
            if (clientMessage.indexOf("[.p.]") >= 0) {
                return null;
            }
        } catch (IOException e) {
            log.error(e);
            connectors.remove(t.getKey());
            sk.cancel();
            if (sk.channel() != null) {
                try {
                    socketChannel.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
            }
            WSocket.getListener().onClose(t.getKey());
            t = null;
        }

        return clientMessage;
    }

    /**
     * send the response to client socket
     * 
     * @throws IOException
     */
    public static void sendResponse(String line, SocketChannel sc)
        throws IOException {
        if (!sc.isOpen()) {
            return;
        }
        byte[] b = new byte[1];
        b[0] = (byte) 0x00;
        byte[] e = new byte[1];
        e[0] = (byte) 0xff;

        sc.write(ByteBuffer.wrap(b));
        sc.write(ByteBuffer.wrap(line.getBytes(CHARSET)));
        sc.write(ByteBuffer.wrap(e)); 
    }

    public static Map<String,String> processHandshake(SocketChannel sc) throws IOException {
        Socket aClientSocket = sc.socket();
        PrintStream os = new PrintStream(aClientSocket.getOutputStream(), true,
        		CHARSET);
        InputStream lIn = aClientSocket.getInputStream();
        byte lBuff[] = new byte[8192];
        int lRead = lIn.read(lBuff);
        if (lRead <= 0) {
            return null;
        }
        byte lReq[] = new byte[lRead];
        System.arraycopy(lBuff, 0, lReq, 0, lRead);

        Map lRespMap = parse2Map(lReq);
        /*
        String lFlashBridgeReq = (String) lRespMap.get("policy-file-request");
        if (lFlashBridgeReq != null) {
        	log.warn((new StringBuilder())
        			.append("Server returned policy file request ('")
        			.append(lFlashBridgeReq)
        			.append("'), check for FlashBridge plug-in.")
        			.toString());
        	byte lBA[] = generateS2CResponse(lRespMap);
            if (lBA == null) {
                return null;
            }
            os.write(lBA);
            os.flush();
            sc.close();
            return null;
        }
        */
        Map<String, String> lArgs = new HashMap<String,String>();
        String lPath = (String) lRespMap.get("path");
        String lSearchString = "";
        if (lPath != null) {
            int lPos = lPath.indexOf(PATHARG_SEPARATOR);
            if (lPos >= 0) {
                lSearchString = lPath.substring(lPos + 1);
                if (lSearchString.length() > 0) {
                    String lArgsArray[] = lSearchString.split(ARGARG_SEPARATOR);
                    for (int i = 0; i < lArgsArray.length; i++) {
                        String lKeyValuePair[] = lArgsArray[i].split(
                        		KEYVAL_SEPARATOR, 2);
                        if (lKeyValuePair.length != 2)
                            continue;
                        lArgs.put(lKeyValuePair[0], lKeyValuePair[1]);
                    }
                }
            }
        }
        String authKey = lArgs.get(KEY_PARAM);
        if (!WSocket.getListener().checkKey(authKey)) {
            sc.close();
            return null;
        }   
            
        byte lBA[] = generateResponse(lRespMap);
        if (lBA == null) {
            return null;
        }
        os.write(lBA);
        os.flush();
        
        return lArgs;
        
    }

    private static long calcSecKeyNum(String aKey) {
        StringBuilder lSB = new StringBuilder();

        int lSpaces = 0;
        for (int i = 0; i < aKey.length(); ++i) {
          char lC = aKey.charAt(i);
          if (lC == ' ')
            ++lSpaces;
          else if ((lC >= '0') && (lC <= '9')) {
            lSB.append(lC);
          }
        }
        long lRes = -1L;
        if (lSpaces > 0) {
          try {
            lRes = Long.parseLong(lSB.toString()) / lSpaces;
          }
          catch (NumberFormatException ex)
          {
          }
        }

        return lRes;
      }

      public static Map parse2Map(byte[] aReq)
      {
        String lHost = null;
        String lOrigin = null;
        String lLocation = null;
        String lPath = null;
        String lSecKey1 = null;
        String lSecKey2 = null;
        byte[] lSecKey3 = new byte[8];
        Boolean lIsSecure = Boolean.valueOf(false);
        Long lSecNum1 = null;
        Long lSecNum2 = null;
        byte[] lSecKeyResp = new byte[8]; 
        Map lRes = new HashMap();

        int lReqLen = aReq.length;
        String lRequest = "";
        try {
          lRequest = new String(aReq, CHARSET);
        }
        catch (Exception ex)
        {
        }
        /*
        if (lRequest.indexOf("policy-file-request") >= 0) {
          lRes.put("policy-file-request", lRequest);
          return lRes;
        }
         */
        lIsSecure = Boolean.valueOf(lRequest.indexOf("Sec-WebSocket") > 0);

        if (lIsSecure.booleanValue()) {
          lReqLen -= 8;
          for (int i = 0; i < 8; ++i) {
            lSecKey3[i] = aReq[(lReqLen + i)];
          }

        }

        int lPos = lRequest.indexOf("Host:");
        lPos += 6;
        lHost = lRequest.substring(lPos);
        lPos = lHost.indexOf("\r\n");
        lHost = lHost.substring(0, lPos);

        lPos = lRequest.indexOf("Origin:");
        lPos += 8;
        lOrigin = lRequest.substring(lPos);
        lPos = lOrigin.indexOf("\r\n");
        lOrigin = lOrigin.substring(0, lPos);

        lPos = lRequest.indexOf("GET");
        lPos += 4;
        lPath = lRequest.substring(lPos);
        lPos = lPath.indexOf("HTTP");
        lPath = lPath.substring(0, lPos - 1);

        lLocation = new StringBuilder().append("ws://").append(lHost).append(lPath).toString();

        lPos = lRequest.indexOf("Sec-WebSocket-Key1:");
        if (lPos > 0) {
          lPos += 20;
          lSecKey1 = lRequest.substring(lPos);
          lPos = lSecKey1.indexOf("\r\n");
          lSecKey1 = lSecKey1.substring(0, lPos);
          lSecNum1 = Long.valueOf(calcSecKeyNum(lSecKey1));
        }

        lPos = lRequest.indexOf("Sec-WebSocket-Key2:");
        if (lPos > 0) {
          lPos += 20;
          lSecKey2 = lRequest.substring(lPos);
          lPos = lSecKey2.indexOf("\r\n");
          lSecKey2 = lSecKey2.substring(0, lPos);
          lSecNum2 = Long.valueOf(calcSecKeyNum(lSecKey2));
        }

        if ((lSecNum1 != null) && (lSecNum2 != null))
        {
          BigInteger sec1 = new BigInteger(lSecNum1.toString());
          BigInteger sec2 = new BigInteger(lSecNum2.toString());

          byte[] l128Bit = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

          byte[] lTmp = sec1.toByteArray();
          int lIdx = lTmp.length;
          int lCnt = 0;
          while ((lIdx > 0) && (lCnt < 4)) {
            --lIdx;
            ++lCnt;
            l128Bit[(4 - lCnt)] = lTmp[lIdx];
          }

          lTmp = sec2.toByteArray();
          lIdx = lTmp.length;
          lCnt = 0;
          while ((lIdx > 0) && (lCnt < 4)) {
            --lIdx;
            ++lCnt;
            l128Bit[(8 - lCnt)] = lTmp[lIdx];
          }

          lTmp = lSecKey3;
          System.arraycopy(lSecKey3, 0, l128Bit, 8, 8);
          try
          {
            MessageDigest md = MessageDigest.getInstance("MD5");
            lSecKeyResp = md.digest(l128Bit);
          }
          catch (Exception ex)
          {
          }
        }
        lRes.put("path", lPath);
        lRes.put("host", lHost);
        lRes.put("origin", lOrigin);
        lRes.put("location", lLocation);
        lRes.put("secKey1", lSecKey1);
        lRes.put("secKey2", lSecKey2);
        lRes.put("isSecure", lIsSecure);
        lRes.put("secKeyResponse", lSecKeyResp);

        return lRes;
      }

      public static byte[] generateResponse(Map aRequest){
    	/*
        String lPolicyFileRequest = (String)aRequest.get("policy-file-request");
        if (lPolicyFileRequest != null) {
          byte[] lBA;
          try {
            lBA = "<cross-domain-policy><allow-access-from domain=\"*\" to-ports=\"*\" /></cross-domain-policy>\n".getBytes(CHARSET);
          } catch (UnsupportedEncodingException ex) {
            lBA = null;
          }
          return lBA;
        }
		*/
        Boolean lIsSecure = (Boolean)aRequest.get("isSecure");
        String lOrigin = (String)aRequest.get("origin");
        String lLocation = (String)aRequest.get("location");
        String lRes = new StringBuilder()
        .append("HTTP/1.1 101 Web")
        .append((lIsSecure.booleanValue()) ? "" : " ")
        .append("Socket Protocol Handshake\r\n")
        .append("Upgrade: WebSocket\r\n")
        .append("Connection: Upgrade\r\n")
        .append((lIsSecure.booleanValue()) ? "Sec-" : "")
        .append("WebSocket-Origin: ").append(lOrigin)
        .append("\r\n")
        .append((lIsSecure.booleanValue()) ? "Sec-" : "")
        .append("WebSocket-Location: ").append(lLocation)
        .append("\r\n")
        .append("\r\n")
        .toString();
        try
        {
          byte[] lBA = lRes.getBytes(CHARSET);

          if (lIsSecure.booleanValue()) {
            byte[] lSecKey = (byte[])(byte[])aRequest.get("secKeyResponse");
            byte[] lResult = new byte[lBA.length + lSecKey.length];
            System.arraycopy(lBA, 0, lResult, 0, lBA.length);
            System.arraycopy(lSecKey, 0, lResult, lBA.length, lSecKey.length);
            return lResult;
          }
          return lBA;
        } catch (UnsupportedEncodingException ex) {
        }
        return null;
      }

}
